import os
import random
import torch
import torch.backends.cudnn as cudnn
from config import get_config
import importlib
import shutil

def main(userconfig):

    sample_dir = os.path.join(userconfig.sample_root_dir, userconfig.trainer, userconfig.net_name)
    if os.path.exists(sample_dir):
        shutil.rmtree(sample_dir)

    if userconfig.refine:
        print("Refine Mode Open")

    checkpoint_dir = os.path.join(userconfig.checkpoint_root_dir, userconfig.trainer, userconfig.net_name)
    os.makedirs(checkpoint_dir, exist_ok=True)
    os.makedirs(sample_dir, exist_ok=True)

    userconfig.manual_seed = random.randint(1, 10000)
    print("Random Seed: ", userconfig.manual_seed)
    random.seed(userconfig.manual_seed)
    torch.manual_seed(userconfig.manual_seed)

    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(userconfig.manual_seed)

    cudnn.benchmark = True
    train_dataset_module = importlib.import_module("datasets.hair.datasetaug")
    train_dataset = train_dataset_module.HairDatasetAug(data_folder=userconfig.trainBasedir,
                                   imglist=userconfig.trainList,
                                   image_size=(userconfig.input_shape[1], userconfig.input_shape[0]),
                                   mask_size=(userconfig.output_shape[1], userconfig.output_shape[0]),
                                   blurmask_size=(userconfig.output_shape[1], userconfig.output_shape[0]),
                                   trimap_size=(userconfig.output_shape[1], userconfig.output_shape[0]),
                                   alpha_size=(userconfig.output_shape[1], userconfig.output_shape[0]),
                                   return_blurmask=userconfig.return_blurmask,
                                   return_mask=userconfig.return_mask,
                                   return_trimap=userconfig.return_trimap,
                                   return_alpha=userconfig.return_alpha)

    train_dataloader = torch.utils.data.DataLoader(dataset=train_dataset,
                                                   batch_size=userconfig.batch_size,
                                                   shuffle=True,
                                                   drop_last=True,
                                                   num_workers=userconfig.num_workers)

    test_dataset_module = importlib.import_module("datasets.hair.dataset")
    test_dataset = test_dataset_module.HairDataset(data_folder=userconfig.testBasedir,
                               imglist=userconfig.testList,
                               image_size=(userconfig.input_shape[1], userconfig.input_shape[0]),
                               mask_size=(userconfig.output_shape[1], userconfig.output_shape[0]),
                               blurmask_size=(userconfig.output_shape[1], userconfig.output_shape[0]),
                               trimap_size=(userconfig.output_shape[1], userconfig.output_shape[0]),
                               alpha_size=(userconfig.output_shape[1], userconfig.output_shape[0]),
                               return_blurmask=userconfig.return_blurmask,
                               return_mask=userconfig.return_mask,
                               return_trimap=userconfig.return_trimap,
                               return_alpha=userconfig.return_alpha)

    test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset,
                                                  batch_size=userconfig.test_batch_size,
                                                  shuffle=False,
                                                  num_workers=4,
                                                  drop_last=False
                                                  )

    module = importlib.import_module("trainer."+userconfig.trainer)
    trainer = module.Trainer(userconfig, train_dataloader, test_dataloader, checkpoint_dir, sample_dir)
    trainer.train()


if __name__ == "__main__":
    userconfig = get_config()
    main(userconfig)
